import networkx as nx
import matplotlib.pyplot as plt

def get_node_name(node):
    if (node.left is None) & (node.right is None):
        return node.val
    return f"{node.id}({node.val})"

def create_graph(G, node, pos={}, color=[], x=0, y=0, level=1):
    # node_id = f"({node.id}, {node.val})"
    node_id = node.get_node_name()
    pos[node_id] = (x, y)
    if not((node.left is None) & (node.right is None)):
        color.append("#00868B")
    else:
        color.append("#FF3030")
    # color.append('r')
    
    if node.left:
        G.add_edge(node_id, node.left.get_node_name())
        l_x, l_y = x - 1 / 2 ** level, y - 1
        l_level = level  + 1
        create_graph(G, node.left, x=l_x, y=l_y, pos=pos, color=color, level=l_level)
    if node.right:
        G.add_edge(node_id, node.right.get_node_name())
        r_x, r_y = x + 1 / 2 ** level, y - 1
        r_level = level + 1
        create_graph(G, node.right, x=r_x, y=r_y, pos=pos, color=color, level=r_level)
    return (G, pos, color)

def draw_tree(root):
    graph = nx.DiGraph()
    graph, pos, color= create_graph(graph, root)
    assert len(pos) == len(color), "pos 长度不等于 color长度"
    fig, ax = plt.subplots(figsize=(16, 10))
    nx.draw_networkx(graph, pos, node_color=color, ax=ax, node_size=1400)
    plt.show()